from Linac import *
from mpi4py import MPI
import storage
from progressbar import Bar, Percentage, RotatingMarker,ETA, ProgressBar

LinacDesign = 'optimusPlus/'
fileDir = '/home/ainswort/HOMDynamics/LatticeFiles/'+LinacDesign
saveDir = '/data/ainswort/Output/'+LinacDesign+'HOMs/Run1/'
fileName = 'Sundelin.h5'
InjFile = '/home/ainswort/HOMDynamics/RunScripts/dtlErrors/'+LinacDesign+'InputData.dat'

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

#####Build linac#####
initialE = 89.6e6

Cavities = CavitySection(file=fileDir+'TW_Sync.txt',E_in=initialE)
Cavities.cavPeriod=np.concatenate([np.ones(26)*2,np.ones(36)*4,np.ones(84)*4])
Cavities.VoltageCorrect()

#####Create Pulse#####
numBunches = int(1e6)
tb = 1/352.2e6
Ib = 62.5e-3
Ibsig = Ib*0.01 

#####Create HOMs#####
homQex = 1e8
fsigS = 10e3
fsigE = 1e3


process = 1
N_tot = 1000
if rank==0:
    pbar  = ProgressBar(widgets=['Linacs Tracked: ', Percentage(), ' ', Bar(marker=RotatingMarker()),' ', ETA()],maxval=N_tot).start()

    fileOut = storage.LinacDB(saveDir+fileName,int(numBunches),'w')
    PSOUT = storage.createPhaseSpaceOutDB(numBunches)
    PSOUTtable = fileOut.h5file.createTable("/", 'PhaseSpaceExit', PSOUT, expectedrows=N_tot, filters=fileOut.filters)

    N = 0

    while process < size and N<N_tot:
        comm.send(N,dest=process,tag=3)
        N+=1
        pbar.update(N)
        process+=1
        
    recieved_processes=0
    while recieved_processes<N_tot:
        L=comm.recv(source=MPI.ANY_SOURCE,tag=1)
        fileOut.Write(L,PSOUTtable,Full=True)
        process = comm.recv(source=MPI.ANY_SOURCE,tag=2) 
        recieved_processes+=1

        if N<N_tot:
            comm.send(dest=process,tag=3)
            N+=1
            pbar.update(N)

    pbar.finish()
    for process in range(1,size):
        comm.send(-1,dest=process,tag=3)


else:
    while True:
        n = comm.recv(source=0,tag=3)
        if n== -1: break


        P1 = Pulse(numBunches,Cavities.cavNum, tb, Ib)
        P1.loadInjectionNoise(InjFile)
        P1.AddCurrentError(Ibsig,seed=4612)
        
        H1 = HOM(Cavities.cavNum,Qex=homQex,nmodes=1)
        H1.loadHomFREQandRQ('/data/ainswort/SPOKE/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([3]),0,26)
        H1.loadHomFREQandRQ('/data/ainswort/ESS067/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([29]),26,36)
        H1.loadHomFREQandRQ('/data/ainswort/ESS086/'+LinacDesign+'Monopoles-RoverQ.dat',np.array([19]),26+36,84)
        H1.RQCorrect(Cavities.W_i,Cavities.cavPeriod)
        H1.AddFrequencyErrorSundelin(352.21e6,0,0,26)
        H1.AddFrequencyErrorSundelin(704.42e6,0,26)
        Lin = Linac(Cavities, P1, H1)
        Lin.Track()

        comm.send(Lin,dest=0,tag=1)
        comm.send(rank,dest=0,tag=2)

